package problems;

import globals.CentralStatics;
import globals.Tuple;
import problems.rdmthreeobjectives.RDMActionsThree;

import java.text.DecimalFormat;
import java.text.NumberFormat;
import java.util.Random;

import value.AlphaMatrix;


public class Mopomdp<S,A,O> {

	protected String problemName;
	
	public static final Integer TERMINAL_STATE = new Integer(-1);
	public static final Integer UNSPECIFIED_STATE = new Integer(-2);
	
	protected boolean finiteHorizon;
	protected int horizon;
	protected double discountFactor;
	
	public int currentState;
	protected double[] bZero;
	public double[] currentBelief;
	
	protected EnumerableStates<S> states;
	protected TransitionFunction<S,A> transitions; 
	protected Actions<A> actions;
	protected RewardFunction<S,A> rewards;
	protected ObservationFunction<S,A,O> observationFunction;
	protected Observations<O> omega;
	
	protected Mopomdp(){
		
	}
	
	public Mopomdp(EnumerableStates<S> s_, TransitionFunction<S,A> t_, Actions<A> a_,
			RewardFunction<S,A> r_, ObservationFunction<S,A,O> o_, Observations<O> om_, 
			double[] b0){
		states = s_;
		transitions = t_;
		actions = a_;
		rewards = r_;
		observationFunction = o_;
		omega = om_;
		
		bZero = b0;
		if(b0 != null){
			this.reInitialize();
		} else {
			System.err.print("No specified initial belief. ");
			System.err.print("Please use appropriate constructor. ");
			System.err.println();
			System.exit(-2);
		}
		
		this.problemName = ""+CentralStatics.getNextProblemIdentifier()+": ";
		
		this.finiteHorizon = false;
		this.horizon = Integer.MAX_VALUE;
		this.discountFactor = 0.9;
	}
	
	public Mopomdp(EnumerableStates<S> s_, TransitionFunction<S,A> t_, Actions<A> a_,
			RewardFunction<S,A> r_, ObservationFunction<S,A,O> o_, Observations<O> om_){
		this(s_,  t_,  a_, r_, o_, om_, Mopomdp.uniformBelief(s_));
	}
	
	public void reInitialize(){
		Random rand = CentralStatics.getCentralRandom();
		currentBelief = this.bZero;
		double stateDbl = rand.nextDouble();
		double curr = 0;
		this.currentState = -1;
		for(int i=0; i<bZero.length; i++){
			curr+=bZero[i];
			if(stateDbl<=curr){
				this.currentState = i;
				break;
			}
			if(i==(bZero.length-1) && this.currentState==-1){
				//Fix for precision errors of doubles; 
				//beliefs might not exactly sum to 1;
				this.currentState = i;
			}
		}
		
		if(currentState==-1){
			System.err.println("(re)initialization of MOPOMDP failed");
			System.exit(-1);
		}
		
	}
	
	//Useful if no initial belief is specified
	public static <S> double[] uniformBelief(EnumerableStates<S> s){
		double[] result = new double[s.numberOfStates()];
		double p = 1.0/((double) result.length);
		for(int i=0; i<result.length; i++){
			result[i]=p;
		}
		return result;
	}
	
	public S currentState(){
		return this.states.stateIdentifier(this.currentState);
	}
	
	public int currentStateNumber(){
		return (this.currentState);
	}
	
	public Tuple<O,double[]> performAction(A action){
		if(currentState==TERMINAL_STATE){
			return null;
		}
		if(currentState==UNSPECIFIED_STATE){
			System.err.println("Trying to simulate action in unspecified state.");
			System.exit(UNSPECIFIED_STATE);
			return null;
		}
		S currentS  = states.stateIdentifier(currentState);
		
		S nextState = this.transitions.nextState(currentS, action);
		
		this.currentState = states.stateNumber(nextState);
		
		
		double[] reward = this.rewards.getReward(currentS, action, nextState);
		
		
		/*if(this.inTerminalState()){
			O obs = this.observationFunction.getObservation(action, nextState);
			//this.reInitialize();
			return new Tuple<O,double[]>(obs,reward);
		} else {*/
			O obs = this.observationFunction.getObservation(action, nextState);
			
			this.beliefUpdate(action, obs);
			return new Tuple<O,double[]>(obs,reward);
		//}
	}
	
	public double[] reward(S state, A action){
		return this.rewards.expectedReward(state, action);
	}
	
	public double[] reward(int state, A action){
		S s = this.states.stateIdentifier(state);
		return this.rewards.expectedReward(s, action);
	}
	
	public Actions<A> getActionSet(){
		return this.actions;
	}
	
	public Observations<O> getObservationSet(){
		return this.omega;
	}
	
	public EnumerableStates<S> getStateSet(){
		return this.states;
	}
	
	private void beliefUpdate(A action, O observation){
		this.currentBelief = this.beliefUpdate(this.currentBelief, action, observation);
	}
	
	public double[] currentBelief(){
		return this.currentBelief;
	}
	
	public double[] beliefUpdate(double[] belief, A action, O observation){
		// b^{a,o}(s') = [O(a,s'o) sum_s T(s,a,s') b(s)] / Pr(o | b,a)
		// Pr(o | b,a) = sum_s b(s) sum_s' T(s,a,s') O(a,s'o)
		
		int numStates= this.states.numberOfStates();
		double[] result = new double[numStates];
		//Calculate Pr(o|b,a)
		double below = 0; //Pr(o|b,a)
		for(int s=0; s<numStates; s++){
			S st = this.states.stateIdentifier(s);
			double margSPrime = 0;
			for(int s_=0; s_<numStates; s_++){
				S stP = this.states.stateIdentifier(s_);
				margSPrime += this.transitions.probabilityOfTransition(st, action, stP)*
						this.observationFunction.observationProbability(action, stP, observation);
			}
			below += belief[s]*margSPrime;
		}
		
		//Calculate probability for every possible next state
		for(int s_=0; s_<numStates; s_++){
			S sPrime = this.states.stateIdentifier(s_);
			double above = 0.0; //O(a,s'o) sum_s T(s,a,s') b(s)
			//Calculate O(a,s'o) sum_s T(s,a,s') b(s)
			for(int s=0; s<numStates; s++){
				S sPrev = this.states.stateIdentifier(s);
				double t = this.transitions.probabilityOfTransition(sPrev, action, sPrime);
				double b = belief[s];
				above += t*b;
			}
			above *= this.observationFunction.observationProbability(action, sPrime, observation);
			result[s_] = above/below;
		}
		
		return result;
	}
	
	public A randomAction(){
		return this.actions.randomAction();
	}
	
	public String toString(){
		String result ="MO-POMDP "+this.problemName+"\n";
		int noS = this.states.numberOfStates();
		result+= ""+noS+" states; currently in state "+this.currentState+"\n";
		result+= "Current belief: [";
		NumberFormat formatter = new DecimalFormat("#0.000");
		for(int i=0; i<noS; i++){
			result+=formatter.format(this.currentBelief[i]);
			if( i != (noS-1) ){
				result+=", ";
			}
		}
		result+="]\n";
		return result;
	}
	
	public int numberOfObjectives(){
		return this.rewards.numberOfObjectives();
	}
	
	public int numberOfStates(){
		return this.states.numberOfStates();
	}
	
	public boolean inTerminalState(){
		//System.out.println("Current state value:"+this.currentState);
		
		return this.states.isTerminal(this.currentState);
	}
	
	public double observationProb(A a, S s_, O o){
		return this.observationFunction.observationProbability(a, s_, o);
	}
	
	public double transitionProb(S s, A a, S s_){
		return this.transitions.probabilityOfTransition(s, a, s_);
	}
	
	public double getGamma(){
		return this.discountFactor;
	}
	
	public boolean isDiscounted(){
		boolean result = this.discountFactor<1;
		return result;
	}
	
	@SuppressWarnings("unchecked")
	public AlphaMatrix<A> misereMatrix(){
		AlphaMatrix<A> mat = null;
		if(this.actions.isDiscrete()){
			DiscreteActions<A,S> acts = (DiscreteActions<A,S>) actions; 
			mat = new AlphaMatrix<A>(this.states.numberOfStates(), this.numberOfObjectives());
			double[] rew = this.rewards.minRewards();
			double multiple;
			if(this.discountFactor<1){
				multiple = 1.0/(1.0-this.discountFactor);
			} else {
				multiple = this.horizon;
			}
			for(int s=0; s<this.states.numberOfStates(); s++){
				this.currentState = s;
				for(int i=0; i<this.numberOfObjectives(); i++){
					if(this.inTerminalState()){
						mat.setValue(s, i, 0);
						//mat.setValue(s, i, rew[i]*multiple);
						///setting the default action///huma
						//mat.theTag=(A) RDMActionsThree.RT;
					} else {
						mat.setValue(s, i, rew[i]*multiple);
						//setting the default action///huma
						mat.theTag=(A) RDMActionsThree.RT;
						
					}
				}
			}
		}
		return mat;
	}
}
